import math
import torch
import torch.nn as nn
import torch.optim as optim

# 参数设置
N = 8  # 天线数量
theta = 30  # 目标角度（度）
c = 3e8  # 光速（m/s）
f = 2.4e9  # 工作频率（Hz）
lambda_ = c / f  # 波长
d = lambda_ / 2  # 天线间距

# 将角度转换为弧度
theta_rad = torch.tensor(math.radians(theta), requires_grad=False)

# 定义天线因子
antenna_factor = torch.tensor([1.0] * N, requires_grad=True)  # 假设所有天线的天线因子相同

# 定义权重系数
weights = torch.randn(N, requires_grad=True)

# 定义波束形成函数
def beamforming(weights, antenna_factor, theta_rad, d, lambda_):
    # 计算每个天线的相位
    phases = weights * 1j
    # 计算每个天线的信号
    signals = antenna_factor * torch.exp(phases)
    # 计算波束方向图
    beam_pattern = torch.abs(torch.sum(signals * torch.exp(-1j * 2 * math.pi * d * torch.arange(N) * torch.sin(theta_rad) / lambda_)))
    return beam_pattern

# 定义波束方向图函数，用于计算整个方向图
def calculate_beam_pattern(weights, antenna_factor, d, lambda_, num_angles=1000):
    angles = torch.linspace(-90, 90, num_angles, requires_grad=False)
    angles_rad = angles / 180.0 * torch.pi
    beam_patterns = []
    for angle_rad in angles_rad:
        pattern = beamforming(weights, antenna_factor, angle_rad, d, lambda_)
        beam_patterns.append(pattern)
    return angles, torch.stack(beam_patterns)

# 定义目标函数
def objective_function(weights, antenna_factor):
    angles, beam_patterns = calculate_beam_pattern(weights, antenna_factor, d, lambda_)
    beam_patterns_db = 20 * torch.log10(beam_patterns)
    max_beam = torch.max(beam_patterns_db)
    # 计算主瓣宽度的平滑近似
    half_power = max_beam - 3  # -3dB
    beam_patterns_smooth = torch.sigmoid(beam_patterns_db - half_power)
    mainlobe_width = torch.sum(beam_patterns_smooth) * (angles[1] - angles[0])
    # 计算旁瓣水平的平滑近似
    sidelobe_level = torch.max(beam_patterns_db[beam_patterns_db < max_beam])
    # 定义权重因子
    alpha = 1.0
    beta = 1.0
    # 计算目标函数
    J = alpha * mainlobe_width + beta * sidelobe_level
    return J

# 定义优化器
optimizer = optim.Adam([weights, antenna_factor], lr=0.01)

# 优化过程
for epoch in range(1000):
    optimizer.zero_grad()
    J = objective_function(weights, antenna_factor)
    J.backward()
    optimizer.step()
    # 确保天线因子非负
    with torch.no_grad():
        antenna_factor.clamp_(min=0.0, max=1.0)
    
    if epoch % 100 == 0:
        print(f'Epoch {epoch}, Objective Function: {J.item()}')

# 打印优化后的权重系数和天线因子
print(f'Optimized Weights: {weights.detach().numpy()}')
print(f'Optimized Antenna Factor: {antenna_factor.detach().numpy()}')
